r'''
把E:\uav_datasets\Drone-Detection-RTB下的对应类别所有图片按照mat里的跟踪框格式重组为16帧一个的视频
'''
import os

import cv2
import scipy.io as sio

src_path = r'E:\uav_datasets\Drone-Detection-RTB\V_AIRPLANE_001'
listdir = os.listdir(src_path)  # 目录下的所有文件，包括mat文件和所有视频帧
# print(len(listdir))
# print(listdir)
# 寻找mat文件
mat_idx, mat_file = -1, ''
for i, f in enumerate(listdir):
    if f.endswith(".mat"):
        mat_file, mat_idx = f, i
        break
assert mat_idx != -1, "未找到mat文件"
vid_frame_names = list(listdir)
vid_frame_names.pop(mat_idx)  # 剔除mat文件
vid_frame_paths = [os.path.join(src_path, x) for x in vid_frame_names]
# print(len(listdir))
listdir = sorted(listdir)  # 按帧编号从小到大排序
# 加载mat文件到内存
mat_path = os.path.join(src_path, mat_file)
# print(mat_path)
loadmat = sio.loadmat(mat_path)
# print(loadmat.keys())
bounding_box = loadmat['bounding_box']
region_track_bbox = loadmat['region_track_bbox']
stem = str(loadmat['stem'][0])
src_vid = str(loadmat['src_path'][0])


def testNdarray():
    # print(bounding_box)
    print(type(bounding_box))  # <class 'numpy.ndarray'>
    print(bounding_box.shape)  # (1, 327)
    print(bounding_box[0].shape)  # (327,)
    gts = bounding_box[0]
    gts = gts[0:3]
    print(gts)  # [313 205  36  24]
    print(gts.shape)  # (4,)
    print("==========")
    # print(type(bounding_box))  # <class 'numpy.ndarray'>
    # print(bounding_box.shape)  # (1, 327)
    # print(bounding_box[0].shape)  # (327,)
    # gts = bounding_box[0]
    # gts = gts[0][0]
    # print(gts)  # [313 205  36  24]
    # print(gts.shape)  # (4,)
    # print("==========")
    # print(type(bounding_box))  # <class 'numpy.ndarray'>
    # print(bounding_box.shape)  # (1, 327)
    # print(bounding_box[0].shape)  # (327,)
    # gts = bounding_box[0]
    # gt0 = gts[0]
    # print(gt0)  # [[313 205  36  24]]
    # print(gt0.shape)  # (1, 4)
    # print("==========")
    # gt0 = gt0[0]
    # print(gt0)  # [313 205  36  24]
    # print(gt0.shape)  # (4,)
    # print("==========")
    # n1, n2 = gt0[0], gt0[1]
    # print(n1, n2)
    # print(n1 > n2)
    # print(n1 + n2)
    # print("==========")
    # print(type(n1)) # <class 'numpy.float64'>
    # print(region_track_bbox)
    # print(stem)


def testDrawRTB():
    cap = cv2.VideoCapture(src_vid)  # 源视频路径
    out_video_path = os.path.join(src_path, stem + ".mp4")  # 样本视频路径
    clip_out = os.path.join(src_path, 'clips')
    if not os.path.exists(clip_out): os.mkdir(clip_out)  # clip视频父文件夹路径

    width = int(cap.get(cv2.CAP_PROP_FRAME_WIDTH))  # 获取原视频的宽
    height = int(cap.get(cv2.CAP_PROP_FRAME_HEIGHT))  # 获取原视频的高
    fps = int(cap.get(cv2.CAP_PROP_FPS))  # 帧率
    fourcc = int(cap.get(cv2.CAP_PROP_FOURCC))  # 维持原视频的编码

    # 视频对象的输出路径及文件名
    print('out_video_path',
          out_video_path)  # out_video_path E:\uav_datasets\Drone-Detection-RTB\V_AIRPLANE_001\V_AIRPLANE_001.mp4
    out = cv2.VideoWriter(out_video_path, fourcc, fps, (width, height))
    # print(region_track_bbox)
    # print(vid_frame_paths[0]) # E:\uav_datasets\Drone-Detection-RTB\V_AIRPLANE_001\V_AIRPLANE_0011_001.png

    i, frame_cnt = 0, 0
    clip_out_writer, rtb = None, None
    for j, vid_frame_path in enumerate(vid_frame_paths):
        if j % 16 == 0:  # 控制region_track_bbox遍历
            if i >= len(region_track_bbox):  # len(region_track_bbox) = 原视频总帧数//16，如327//16=20
                if clip_out_writer: clip_out_writer.release()
                break
            rtb = [int(x) for x in region_track_bbox[i]]  # 获取第i个rtb
            p1, p2 = (rtb[0], rtb[1]), (rtb[2], rtb[3])
            clip_out_path = os.path.join(clip_out, stem + "_" + str(j) + ".mp4")  # clip存放路径：src_path/clips/xx_x.mp4
            clip_h, clip_w = rtb[2] - rtb[0], rtb[3] - rtb[1]  # clip的高和宽
            print('=====================clip_h, clip_w', clip_h, clip_w)
            if clip_out_writer: clip_out_writer.release()
            clip_out_writer = cv2.VideoWriter(clip_out_path, fourcc, fps, (clip_h, clip_w))
            i += 1
        vid_frame = cv2.imread(vid_frame_path)
        clip_frame = vid_frame[rtb[1]:rtb[3], rtb[0]:rtb[2]]  # 裁取clip区域
        print('clip_frame.shape', clip_frame.shape)
        clip_out_writer.write(clip_frame)  # 往单个clip里添加视频帧

        cv2.rectangle(vid_frame, p2, p1, (0, 0, 128), 1)
        out.write(vid_frame)  # 往sample视频里添加视频帧
        frame_cnt += 1
    cap.release()
    out.release()
    print(frame_cnt)  # 正常=320


# 加载原视频信息（如帧速率高度宽度等），保存到新视频
def testVidLoadAndSave():
    print(src_vid)
    out_video_path = os.path.join(src_path, stem + ".mp4")
    cap = cv2.VideoCapture(src_vid)
    width = int(cap.get(cv2.CAP_PROP_FRAME_WIDTH))  # 获取原视频的宽
    height = int(cap.get(cv2.CAP_PROP_FRAME_HEIGHT))  # 获取原视频的高
    fps = int(cap.get(cv2.CAP_PROP_FPS))  # 帧率
    fourcc = int(cap.get(cv2.CAP_PROP_FOURCC))  # 维持原视频的编码
    # 视频对象的输出路径及文件名
    print('out_video_path', out_video_path)
    out = cv2.VideoWriter(out_video_path, fourcc, fps, (width, height))
    retaining, frame = cap.read()
    while retaining:
        out.write(frame)
        retaining, frame = cap.read()
    cap.release()
    out.release()


def testSplit():
    print(src_vid)
    print(type(src_vid))
    src_vid_s = str(src_vid[0])
    print(src_vid_s)
    print(type(src_vid_s))

    # print(region_track_bbox)


if __name__ == '__main__':
    # testNdarray()
    # testSplit()
    # testVidLoadAndSave()
    # testDrawRTB()
    l1 = ['5', '4332', '6', '7']
    l1.sort()
    print(l1)
